#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ROS 节点：doublo_yolo_node_0816

功能：
- 订阅 /yolo_detect(String)。收到 "yolostart" 后，从广角 GS 取帧，运行双阶段 YOLO（double_yolo_pipeline）流程；
- 每帧打印识别结果，并将非 None 的结果分别写入颜色/数字/仪表盘三个缓存队列；
- 当三个队列均满时，进行多数投票结算，按 "color-number-state" 组合，发布到 /yolo_result(String)；
- 发布后暂停检测循环与管线，不再从 GS 取帧，但不 kill 掉 GS（保持 GS 线程存活）。
- 节点启动时会完成一次 warm up：初始化 GS 并获取至少一帧，初始化 YOLO 并跑一次推理，以便 yolostart 后可立即识别。
"""

import os
import sys
import time
import threading
from collections import deque, Counter
from typing import Optional, Dict, Any

import rospy
from std_msgs.msg import String


def _add_import_paths():
    """将 yolo-jetson 相关目录加入 sys.path 以便导入。"""
    current_dir = os.path.dirname(os.path.abspath(__file__))
    project_root = os.path.abspath(os.path.join(current_dir, "../../../../"))
    detect_dir = os.path.join(project_root, "yolo-jetson", "detect")
    rs2_dir = os.path.join(project_root, "yolo-jetson", "rs2")
    for p in [detect_dir, rs2_dir]:
        if os.path.exists(p) and p not in sys.path:
            sys.path.insert(0, p)


_add_import_paths()

try:
    from double_yolo_pipeline import DoubleYoloPipeline
except Exception as e:
    rospy.logerr(f"导入 DoubleYoloPipeline 失败: {e}")
    raise

try:
    from GStreamerWrapper import GStreamerWrapper
except Exception as e:
    rospy.logerr(f"导入 GStreamerWrapper 失败: {e}")
    raise


class DoubloYoloNode0816:
    def __init__(self) -> None:
        rospy.init_node("doublo_yolo_node_0816", anonymous=True)

        # 参数
        self.model_path: str = rospy.get_param(
            "~model_path",
            os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))),
                         "yolo-jetson", "weights", "all_yolo_all_epochs_0811", "best.pt"),
        )
        self.device: Optional[str] = rospy.get_param("~device", None)
        self.debug: bool = bool(rospy.get_param("~debug", True))
        self.save_dir: str = rospy.get_param("~save_dir", "/tmp/doublo_yolo_0816")

        # GS 取流参数
        self.gs_width: int = int(rospy.get_param("~gs_width", 640))
        self.gs_height: int = int(rospy.get_param("~gs_height", 360))
        self.frame_poll_interval: float = float(rospy.get_param("~frame_poll_interval", 0.02))

        # 队列长度（三个队列都需要达到此长度才结算）
        self.queue_length: int = int(rospy.get_param("~queue_length", 7))

        # 双阶段 YOLO 配置
        self.top_cut_ratio: float = float(rospy.get_param("~top_cut_ratio", 0.25))
        self.width_center_ratio: float = float(rospy.get_param("~width_center_ratio", 0.5))
        self.first_stage_conf: float = float(rospy.get_param("~first_stage_conf", 0.6))
        self.second_stage_conf: float = float(rospy.get_param("~second_stage_conf", 0.5))
        self.final_conf: float = float(rospy.get_param("~final_conf", 0.6))

        # 状态
        self.is_running: bool = False
        self.stop_event: threading.Event = threading.Event()
        self.det_thread: Optional[threading.Thread] = None
        self.frame_idx: int = 0
        self.total_processed: int = 0

        # 结果缓存队列
        self.color_queue: deque = deque(maxlen=self.queue_length)
        self.number_queue: deque = deque(maxlen=self.queue_length)
        self.dashboard_queue: deque = deque(maxlen=self.queue_length)

        # 发布者与订阅者（注意：发布话题为 /yolo_resul 按需求保持原样拼写）
        self.result_pub = rospy.Publisher("/yolo_result", String, queue_size=10)
        self.cmd_sub = rospy.Subscriber("/yolo_detect", String, self._command_cb, queue_size=1)

        # GS 与 YOLO 管道
        self.gs: Optional[GStreamerWrapper] = None
        self.pipeline: Optional[DoubleYoloPipeline] = None

        rospy.loginfo("[doublo_yolo_node_0816] 初始化中...")
        # 初始化 GS 并做一次 warm up
        self._warmup_once()
        rospy.loginfo("[doublo_yolo_node_0816] 初始化完成，等待 'yolostart' 指令")

    # ---------- 生命周期 / 控制 ----------
    def _command_cb(self, msg: String) -> None:
        cmd = (msg.data or "").strip().lower()
        rospy.loginfo(f"收到命令: {cmd}")
        if cmd == "yolostart":
            self.start()
        elif cmd == "yolostop":
            self.stop()
        else:
            rospy.logwarn(f"未知命令: {cmd}")

    def start(self) -> None:
        if self.is_running:
            rospy.loginfo("检测已在运行中")
            return
        # 清队列
        self.color_queue.clear()
        self.number_queue.clear()
        self.dashboard_queue.clear()
        # 状态复位
        self.is_running = True
        self.stop_event.clear()
        self.frame_idx = 0
        self.total_processed = 0
        # 启动线程
        self.det_thread = threading.Thread(target=self._loop, name="doublo_yolo_loop", daemon=True)
        self.det_thread.start()
        rospy.loginfo("✅ 检测已启动")

    def stop(self) -> None:
        self.is_running = False
        self.stop_event.set()
        if self.det_thread and self.det_thread.is_alive():
            try:
                self.det_thread.join(timeout=2.0)
            except Exception:
                pass
        self.det_thread = None
        # 保持 GS 存活，不调用 StopThread（不 kill GS）
        rospy.loginfo("🛑 检测已暂停（GS 保持存活）")

    # ---------- 预热 ----------
    def _get_one_frame(self, timeout_s: float = 2.0) -> Optional[Any]:
        deadline = time.time() + float(timeout_s)
        if self.gs is None:
            try:
                self.gs = GStreamerWrapper(width=int(self.gs_width), height=int(self.gs_height))
            except Exception as e:
                rospy.logwarn(f"初始化 GStreamerWrapper 失败: {e}")
                return None
        while time.time() < deadline:
            frame = self.gs.GetFrame()
            if frame is not None:
                return frame
            time.sleep(0.01)
        return None

    def _warmup_once(self) -> None:
        # 初始化 YOLO 管道
        if self.pipeline is None:
            if not os.path.exists(self.model_path):
                rospy.logwarn(f"模型文件不存在，预热跳过: {self.model_path}")
                return
            self.pipeline = DoubleYoloPipeline(
                model_path=self.model_path,
                device=self.device,
                debug=self.debug,
                save_dir=self.save_dir,
            )
        # 获取一帧并跑一次推理
        frame = self._get_one_frame(timeout_s=2.0)
        if frame is None:
            import numpy as np
            frame = np.zeros((int(self.gs_height), int(self.gs_width), 3), dtype=np.uint8)
        try:
            _ = self.pipeline.process_image(
                image=frame,
                top_cut_ratio=self.top_cut_ratio,
                width_center_ratio=self.width_center_ratio,
                first_stage_conf=self.first_stage_conf,
                second_stage_conf=self.second_stage_conf,
                final_conf=self.final_conf,
                frame_idx=0,
            )
            rospy.loginfo("Warm up 完成")
        except Exception as e:
            rospy.logwarn(f"Warm up 推理失败: {e}")

    # ---------- 检测主循环 ----------
    def _value_ok(self, v: Optional[str]) -> bool:
        if v is None:
            return False
        s = str(v).strip().lower()
        return s not in ("", "none", "unknown", "nan")

    def _majority(self, dq: deque) -> str:
        if not dq:
            return "None"
        c = Counter(dq)
        max_cnt = max(c.values())
        cand = [k for k, v in c.items() if v == max_cnt]
        if len(cand) == 1:
            return cand[0]
        for v in reversed(dq):
            if v in cand:
                return v
        return cand[0]

    def _loop(self) -> None:
        rospy.loginfo("识别线程已启动")
        while self.is_running and not self.stop_event.is_set():
            try:
                if self.gs is None:
                    time.sleep(0.05)
                    continue
                frame = self.gs.GetFrame()
                if frame is None:
                    time.sleep(self.frame_poll_interval)
                    continue
                self.frame_idx += 1
                t0 = time.time()
                result: Dict[str, Any] = self.pipeline.process_image(
                    image=frame,
                    top_cut_ratio=self.top_cut_ratio,
                    width_center_ratio=self.width_center_ratio,
                    first_stage_conf=self.first_stage_conf,
                    second_stage_conf=self.second_stage_conf,
                    final_conf=self.final_conf,
                    frame_idx=self.frame_idx,
                )
                dt = time.time() - t0
                self.total_processed += 1

                cone_color = str(result.get("cone_color", "None"))
                number = str(result.get("number", "None"))
                dashboard = str(result.get("dashboard_state", "None"))

                rospy.loginfo(f"[Frame {self.frame_idx}] {cone_color}-{number}-{dashboard} | {dt:.3f}s")

                if self._value_ok(cone_color):
                    self.color_queue.append(cone_color)
                if self._value_ok(number):
                    self.number_queue.append(number)
                if self._value_ok(dashboard):
                    self.dashboard_queue.append(dashboard)

                # 队列长度日志
                rospy.loginfo(
                    f"Q 状态: color={len(self.color_queue)}/{self.queue_length}, "
                    f"number={len(self.number_queue)}/{self.queue_length}, "
                    f"state={len(self.dashboard_queue)}/{self.queue_length}"
                )

                # 三个队列都满 -> 结算并发布
                if (len(self.color_queue) == self.queue_length and
                    len(self.number_queue) == self.queue_length and
                    len(self.dashboard_queue) == self.queue_length):
                    final_color = self._majority(self.color_queue)
                    final_number = self._majority(self.number_queue)
                    final_state = self._majority(self.dashboard_queue)
                    final_str = f"{final_color}-{final_number}-{final_state}"

                    msg = String()
                    msg.data = final_str
                    self.result_pub.publish(msg)
                    rospy.loginfo(f"✅ 最终结果: {final_str} (已发布到 /yolo_result)")

                    # 暂停检测循环，不再从 GS 获取图像（不 kill GS）
                    self.stop()
                    break

                time.sleep(self.frame_poll_interval)
            except Exception as e:
                rospy.logerr(f"检测异常: {e}")
                time.sleep(0.1)

    # ---------- 运行 ----------
    def run(self) -> None:
        rospy.loginfo("doublo_yolo_node_0816 运行中，等待 /yolo_detect 指令")
        rospy.spin()


if __name__ == "__main__":
    try:
        node = DoubloYoloNode0816()
        node.run()
    except rospy.ROSInterruptException:
        pass
    except Exception as e:
        rospy.logerr(f"节点异常退出: {e}")

